"""
Gradio integration for rStar-Math
"""
import os
import json
import gradio as gr
import numpy as np
import pandas as pd
import plotly.express as px
from dotenv import load_dotenv
from src.core.mcts import MCTS
from src.core.ppm import ProcessPreferenceModel
from src.models.model_interface import ModelFactory

# Load environment variables from .env file
load_dotenv()

class RStarMathGradio:
    def __init__(self):
        """Initialize rStar-Math components."""
        self.config_path = 'config/default.json'
        self.mcts = MCTS.from_config_file(self.config_path)
        self.ppm = ProcessPreferenceModel.from_config_file(self.config_path)
        
        # Load available models from config
        self.models = {}
        self.load_available_models()
    
    def load_available_models(self):
        """Load models that have API keys available."""
        with open(self.config_path, 'r') as f:
            config_data = json.load(f)
        
        model_configs = config_data.get('models', {})
        
        for model_name, model_config in model_configs.items():
            try:
                # Try to create model instance using new ModelFactory API
                model = ModelFactory.create_model(
                    model_name=model_name,
                    config_path=self.config_path
                )
                self.models[model_name] = model
                print(f"✅ Successfully loaded model: {model_name}")
            except Exception as e:
                print(f"❌ Failed to load model {model_name}: {str(e)}")
        
        if not self.models:
            print("⚠️ No models loaded. Please check your API key environment variables.")
    
    def solve_problem(self,
                     problem: str,
                     model_name: str,
                     use_rstar: bool = True,
                     show_visualization: bool = True):
        """Solve math problem and return solution with visualizations."""
        if not self.models:
            return "Error: No models available. Please check your API key configuration.", None
            
        if model_name not in self.models:
            return f"Error: Model '{model_name}' not available.", None
            
        model = self.models[model_name]
        
        if use_rstar:
            action, trajectory = self.mcts.search(problem)
            solution_steps = []
            confidence_scores = []
            
            for step in trajectory:
                confidence = self.ppm.evaluate_step(step['state'], model)
                solution_steps.append(step['state'])
                confidence_scores.append(confidence)
        else:
            solution = model.generate_response(problem)
            confidence = model.evaluate_reasoning(problem, [solution])
            solution_steps = [solution]
            confidence_scores = [confidence]
            
        # Format output
        output = "Solution Steps:\n"
        for i, (step, conf) in enumerate(zip(solution_steps, confidence_scores), 1):
            output += f"Step {i}: {step}\nConfidence: {conf:.2f}\n\n"
            
        # Create visualization
        if show_visualization and len(confidence_scores) > 1:
            df = pd.DataFrame({
                'Step': range(1, len(confidence_scores) + 1),
                'Confidence': confidence_scores
            })
            fig = px.line(df, x='Step', y='Confidence',
                         title='Solution Confidence Trend')
            
            return output, fig
        
        return output, None

def create_examples(available_models):
    """Create example problems for the demo."""
    default_model = available_models[0] if available_models else ""
    return [
        ["What is 15 × 27?", default_model, True, True],
        ["Solve for x: 2x + 5 = 13", default_model, True, True],
        ["Find the derivative of f(x) = x² + 3x", default_model, True, True],
        ["Find the area of a circle with radius 5", default_model, True, True]
    ]

def main():
    """Create and launch Gradio interface."""
    rstar = RStarMathGradio()
    
    # Check if any models are available
    if not rstar.models:
        print("❌ No models available. Please configure API keys and try again.")
        print("\nAvailable model configurations:")
        with open('config/default.json', 'r') as f:
            config_data = json.load(f)
        for model_name, config in config_data.get('models', {}).items():
            env_var = config.get('api_key_env_name', f"{model_name.upper()}_API_KEY")
            print(f"  - {model_name}: Set environment variable {env_var}")
        return
    
    # Create interface
    with gr.Blocks(title="rStar-Math Demonstrator") as demo:
        gr.Markdown("# rStar-Math Problem Solver")
        
        with gr.Row():
            with gr.Column():
                problem_input = gr.Textbox(
                    label="Enter your math problem",
                    placeholder="e.g., What is 2 + 2?"
                )
                model_select = gr.Dropdown(
                    choices=list(rstar.models.keys()),
                    value=list(rstar.models.keys())[0] if rstar.models else None,
                    label="Select Model"
                )
                use_rstar = gr.Checkbox(
                    label="Use rStar-Math Enhancement",
                    value=True
                )
                show_viz = gr.Checkbox(
                    label="Show Visualization",
                    value=True
                )
                solve_btn = gr.Button("Solve")
            
            with gr.Column():
                solution_output = gr.Textbox(
                    label="Solution",
                    lines=10
                )
                plot_output = gr.Plot(label="Confidence Trend")
        
        # Add examples
        gr.Examples(
            examples=create_examples(list(rstar.models.keys())),
            inputs=[problem_input, model_select, use_rstar, show_viz]
        )
        
        # Connect components
        solve_btn.click(
            fn=rstar.solve_problem,
            inputs=[problem_input, model_select, use_rstar, show_viz],
            outputs=[solution_output, plot_output]
        )
        
        # Add model status info
        with gr.Accordion("Model Status", open=True):
            model_status = "**Available Models:**\n"
            for model_name in rstar.models.keys():
                model_status += f"✅ {model_name}\n"
            gr.Markdown(model_status)
        
        # Add documentation
        with gr.Accordion("About"):
            gr.Markdown("""
            ## rStar-Math Demonstrator
            
            This interface demonstrates the capabilities of rStar-Math, an AI framework
            that enhances mathematical reasoning using Monte Carlo Tree Search and
            Process Preference Models.
            
            ### Features:
            - Multiple LLM support (OpenAI, Anthropic, Mistral, etc.)
            - Step-by-step solution breakdown
            - Confidence scoring for each step
            - Visual confidence tracking
            
            ### Usage:
            1. Enter your math problem
            2. Select a model
            3. Choose whether to use rStar-Math enhancement
            4. Click "Solve" to see the solution
            """)
    
    # Launch interface
    demo.launch(share=False)

if __name__ == "__main__":
    main()
